from flask import Flask, request, jsonify
from random import Random
import cv2
import onnxruntime as ort
import numpy as np

# Flask app
app = Flask(__name__)

# random number generator
random = Random()
random.seed(0)


def preprocess(image):
    # resize
    resized_image = cv2.resize(image, (320, 320))
    # convert color channel
    rgb_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2RGB)

    # to float && normalize
    normalized_image = rgb_image.astype(np.float32) / 255.0

    # batch dimension
    input_data = np.expand_dims(normalized_image, axis=0)
    # order dimensions
    return input_data.transpose([0, 3, 1, 2])


def process(input_data):
    session = ort.InferenceSession('YoloV7_Tiny.onnx')
    outputs = session.run(['output'], {'images': input_data})[0]
    mapping = [3, 3, 0, 3, 3, 3, 3, 3, 1, 2, 2, 3, 1, 1, 3, 3, 3, 3, 3]
    for output in outputs:
        output[5] = mapping[int(output[5])]
    return outputs


@app.post("/api/image_process")
def image_process():
    data = []
    # get image from request
    if 'image' not in request.files:
        return {"message": "image file required"}, 400
    image_data = request.files['image']
    image = cv2.imdecode(np.frombuffer(image_data.stream.read(), np.uint8), cv2.IMREAD_COLOR)

    # preprocess
    input = preprocess(image)

    # process
    outputs = process(input)
    for output in outputs:
        data.append({
            "box_1": float(output[1]),
            "box_2": float(output[2]),
            "box_3": float(output[3]),
            "box_4": float(output[4]),
            "gesture": int(output[5]),
            "score": float(output[6]),
        })

    # return processed data
    return jsonify(data), 200


def judge(a: int, b: int):
    # 0 - draw, 1 - a wins, 2 - b wins
    result = (a - b + 3) % 3
    if result == 0:
        return "draw"
    elif result == 1:
        return "win"
    else:
        return "lose"


@app.post("/api/judge")
def game_judge():
    if not request.json or 'gesture' not in request.json:
        return {'message': 'gesture number required. '}, 400
    # get player gesture for request
    # 0 - rock, 1 - paper, 2 - scissors
    player_gesture = int(request.json["gesture"])
    opponent_gesture = random.randint(0, 2)
    data = {'result': judge(player_gesture, opponent_gesture), 'opponent': opponent_gesture}
    return jsonify(data), 200


@app.route("/")
def index():
    index_page = ""
    with open("index.html", encoding='utf-8') as f:
        index_page = f.read()
    return index_page, 200


if __name__ == '__main__':
    app.run()
